import os
import cv2
import yaml
import time
import numpy as np
import random

# from TensorRT.DetFM import ToyotaDet
# from TensorRT.DetZM import ToyotaCls
# from TensorRT.Tools import match

import sys
try:
    from DetFM import ToyotaDet
    from DetZM import ToyotaCls
    from Tools import match
    from DetectionZM.XD.XD import xd 
    from DetectionZM.YT.YT import yt
    from DetectionZM.UPR.UPR import upr
    from DetectionZM.M.M import m
    from DetectionZM.LXZST.LXZST import lxzst
    from DetectionZM.LBZ.LBZ import lbz
    from DetectionZM.AR.AR import ar
    from DetectionZM.ZST.ZST import zst


    from DetectionFM.LS.LS import ls
    from DetectionFM.KK.KK import kk
    from DetectionFM.RD.RD import rd
    from DetectionFM.LSH.LSH import lsh
    from DetectionFM.XS.XS import xs
    from DetectionFM.TP.TP import tp
    from DetectionFM.EAPM.EAPM import eapm
    from DetectionFM.BarCode.OCR import ocr_get_number
    from DetectionFM.BarCode.OCR import TextSystem

    from Models.classification.mobilenet import lskkrd_cls
except:
    from .DetFM import ToyotaDet
    from .DetZM import ToyotaCls
    from .Tools import match
    from .DetectionZM.XD.XD import xd 
    from .DetectionZM.YT.YT import yt
    from .DetectionZM.UPR.UPR import upr
    from .DetectionZM.M.M import m
    from .DetectionZM.LXZST.LXZST import lxzst
    from .DetectionZM.LBZ.LBZ import lbz
    from .DetectionZM.AR.AR import ar
    from .DetectionZM.ZST.ZST import zst


    from .DetectionFM.LS.LS import ls
    from .DetectionFM.KK.KK import kk
    from .DetectionFM.RD.RD import rd
    from .DetectionFM.LSH.LSH import lsh
    from .DetectionFM.XS.XS import xs
    from .DetectionFM.TP.TP import tp
    from .DetectionFM.EAPM.EAPM import eapm
    from .DetectionFM.BarCode.OCR import ocr_get_number
    from .DetectionFM.BarCode.OCR import TextSystem


    from .Models.classification.mobilenet import lskkrd_cls

def generateKeyPoint(points):
    keypoints = []
    for i, point in enumerate(points):
        temp = cv2.KeyPoint(x=point[0][0], y=point[0][1], size=point[1], angle=point[2], \
                response=point[3], octave=point[4], class_id=point[5])
        keypoints.append(temp)
    return keypoints

class Toyota(object):
    def __init__(self, yaml_dir):
        self.cfg = yaml.load(open(yaml_dir), Loader=yaml.FullLoader)
        self.algorithm_dir = self.cfg['abs_dir']

        shb_cls, zst_cls, lbz_cls = self.cfg['shb_cls'], self.cfg['zst_cls'], self.cfg['lbz_cls']
        shb_num_classes, zst_num_classes, lbz_num_classes = self.cfg['shb_num_classes'], self.cfg['zst_num_classes'], self.cfg['lbz_num_classes']
        imgsz_cls = self.cfg['imgsz_cls']

        fm_det, zm_det = self.cfg['fm_det'], self.cfg['zm_det']
        imgsz_det = self.cfg['imgsz_det']

        device = self.cfg['device']
        self.template_path = self.cfg['template_path']
        self.plot_det, self.plot_cls = self.cfg['plot_det'], self.cfg['plot_cls']
        # yolo fm
        self.model_fm_det = ToyotaDet(weights=os.path.join(self.algorithm_dir, fm_det), device=device, imgsz=imgsz_det)
        # yolo zm
        self.model_zm_det = ToyotaDet(weights=os.path.join(self.algorithm_dir, zm_det), device=device, imgsz=imgsz_det)
        # color classification
        self.shb_cls = ToyotaCls(weights=os.path.join(self.algorithm_dir, shb_cls), num_classes=shb_num_classes, device=device, imgsz=imgsz_cls)
        self.zst_cls = ToyotaCls(weights=os.path.join(self.algorithm_dir, zst_cls), num_classes=zst_num_classes, device=device, imgsz=imgsz_cls)
        self.lbz_cls = ToyotaCls(weights=os.path.join(self.algorithm_dir, lbz_cls), num_classes=lbz_num_classes, device=device, imgsz=imgsz_cls)

        # rd ls kk init
        self.rd_cls = lskkrd_cls(n_class=2, weights=os.path.join(self.algorithm_dir,self.cfg['rd_model_cls']), device=device)
        self.kk_cls = lskkrd_cls(n_class=2, weights=os.path.join(self.algorithm_dir, self.cfg['kk_model_cls']), device=device)
        self.ls_cls = lskkrd_cls(n_class=2, weights=os.path.join(self.algorithm_dir, self.cfg['ls_model_cls']), device=device)
        
        # OCR init
        self.text_sys = TextSystem(os.path.join(self.algorithm_dir, self.cfg['ocr_det_model_path']),
                          os.path.join(self.algorithm_dir, self.cfg['ocr_rec_model_path']),
                          text_score=self.cfg['ocr_score'],
                          use_angle_cls=True,
                          cls_model_path=os.path.join(self.algorithm_dir, self.cfg['ocr_cls_model_path']),
                          keys_path=os.path.join(self.algorithm_dir, self.cfg['ocr_keys_path']))

        # 门板类型
        self.code = None

        self.flag_kk, self.flag_ls, self.flag_xs, self.flag_lsh = None, None, None, None
        self.flag_rd, self.flag_tp, self.flag_eapm, self.flag_mz = None, None, None, None
        self.flag_bq, self.flag_mq, self.flag_ewm, self.flag_xd = None, None, None, None

        self.flag_zst, self.flag_upr, self.flag_m = None, None, None
        self.flag_ar, self.flag_yt, self.flag_md = None, None, None
        self.flag_lxzst = None

        self.image_ls8, self.image_ls7 = None, None
        self.image_kk8, self.image_kk7 = None, None
        self.image_rd8, self.image_rd7 = None, None

    def get_imgs(self, data):
        """
        self.image1 反面 左
        self.image2 反面 右
        self.image3 正面 左
        self.image4 正面 右
        self.image5 侧

        """
        prefix = "" if self.cfg['image_prefix'] is None else self.cfg['image_prefix']
        self.image1 = data[prefix+'192.168.8.8']
        self.image2 = data[prefix+'192.168.7.7']
        self.image3 = data[prefix+'192.168.4.5']
        self.image4 = data[prefix+'192.168.10.10']
        self.image5 = data[prefix+'192.168.9.9']

    def set_images(self):
        data = dict()
        if self.code is None:
            data['192.168.8.8'] = cv2.cvtColor(self.image1, cv2.COLOR_BGR2RGB)
            data['192.168.7.7'] = cv2.cvtColor(self.image2, cv2.COLOR_BGR2RGB)
            # fm
            data['192.168.4.5'] = cv2.cvtColor(self.image3, cv2.COLOR_RGB2BGR) 
            data['192.168.10.10'] = cv2.cvtColor(self.image4, cv2.COLOR_RGB2BGR) 
            data['192.168.9.9'] = cv2.cvtColor(self.image5, cv2.COLOR_RGB2BGR) 
        else:
            # zm
            data['192.168.8.8'] = cv2.cvtColor(self.image6, cv2.COLOR_BGR2RGB)
            data['192.168.7.7'] = cv2.cvtColor(self.image7, cv2.COLOR_BGR2RGB)
            # fm
            data['192.168.4.5'] = cv2.cvtColor(self.image_zm_left, cv2.COLOR_RGB2BGR) 
            data['192.168.10.10'] = cv2.cvtColor(self.image_zm_right, cv2.COLOR_RGB2BGR) 
            data['192.168.9.9'] = cv2.cvtColor(self.image5, cv2.COLOR_RGB2BGR) 
            # kk\ls\rd result
            if self.flag_tp:
                data['LS192.168.8.8'] = self.image_ls8 
                data['LS192.168.7.7'] = self.image_ls7
                data['RD192.168.8.8'] = self.image_rd8
                data['RD192.168.7.7'] = self.image_rd7
                data['KK192.168.8.8'] = self.image_kk8
                data['KK192.168.7.7'] = self.image_kk7
        return data

    def set_result(self):
        result = dict()
        
        # fm
        result['kk'] = True if self.flag_kk is None else self.flag_kk
        result['ls'] = True if self.flag_ls is None else self.flag_ls
        result['xs'] = True if self.flag_xs is None else self.flag_xs
        result['lsh'] = True if self.flag_lsh is None else self.flag_lsh
        result['rd'] = True if self.flag_rd is None else self.flag_rd
        result['tp'] = True if self.flag_tp is None else self.flag_tp
        result['eapm'] = True if self.flag_eapm is None else self.flag_eapm
        result['bq'] = True if self.flag_bq is None else self.flag_bq
        result['barcode'] = self.code 
        # no detection
        #result['mq'] = True if self.flag_mq is None else self.flag_mq
        #result['ewm'] = True if self.flag_ewm is None else self.flag_ewm
        #result['mz'] = True if self.flag_mz is None else self.flag_mz
        # zm
        result['xd'] = True if self.flag_xd is None else self.flag_xd
        result['zst'] = True if self.flag_zst is None else self.flag_zst
        result['upr'] = True if self.flag_upr is None else self.flag_upr
        result['m'] = True if self.flag_m is None else self.flag_m
        result['ar'] = True if self.flag_ar is None else self.flag_ar
        result['yt'] = True if self.flag_yt is None else self.flag_yt
        result['lxzst'] = True if self.flag_lxzst is None else self.flag_lxzst
        # can not
        result['md'] = True if self.flag_md is None else self.flag_md # 不检测

        
        return result

    def deal_fm_attr(self):
        # yolov_t0 = time.time()
        self.image6, self.Item1 = self.model_fm_det.predict(source=self.image1, plot=self.plot_det)
        self.image7, self.Item2 = self.model_fm_det.predict(source=self.image2, plot=self.plot_det)
        # yolov_t1 = time.time()
        # print("反面目标检测耗时：", yolov_t1-yolov_t0)
        self.det_bq(self.Item1, self.Item2)
        self.flag_bq = True
        if self.flag_bq:
            barcode = self.det_bq_ocr(self.cfg, self.Item1, self.Item2)
            self.code = barcode

            
        else:
            print(5*"*", "门板无标签", 5*"*")
            return None
        # ocr_t2 = time.time()
        if barcode is not None:
            # print("OCR识别门板耗时：", ocr_t2-yolov_t1)
            return barcode
        else:
            print(5*"*","未能识别出标签信息", 5*"*")
            return None

    def deal_fm(self, barcode):

        if barcode is not None:
            self.det_lsh(self.cfg, self.Item1, barcode)
            self.det_xs(self.cfg, self.Item2, barcode)

            # print(Item1, '\n', Item2)
            self.det_tp(self.cfg, self.Item1, self.Item2, barcode)
            
            self.det_eapm(self.cfg, self.Item1, self.Item2, barcode)
            
            self.det_ls(self.cfg, self.Item1, self.Item2, barcode) # cam11, cam13
            self.det_rd(self.cfg, self.Item1, self.Item2, barcode) # cam11, cam13
            self.det_kk(self.cfg, self.Item1, self.Item2, barcode) # cam11, cam13

            self.det_ewm()
            self.det_mq()

    def style(self):

        pass

   
    def deal_zm(self, barcode):
        if barcode is not None:
            self.image_zm_left, self.Item_zm_left = self.model_zm_det.predict(self.image3, plot=self.plot_det) # left
            self.image_zm_right, self.Item_zm_right = self.model_zm_det.predict(self.image4, plot=self.plot_det) # right

            # print(self.Item_zm_left)
            # print(self.Item_zm_right)
            self.det_yt(self.cfg)

            self.det_m(self.cfg, barcode) # 检测上护板
            self.det_lxzst(self.cfg)

            image_zm = {
                'left': self.image3,
                'right': self.image4
            }
            self.det_lbz(self.cfg, image_zm, barcode) # 检测喇叭罩

            image_side = {
                'side': self.image5,
                'desc': 'side camera'
            } 
            self.det_zst(self.cfg, image_side , barcode) # 检测装饰条

            # #start = time.time()
            # attrs = np.load(os.path.join(self.template_path, barcode+'10.npy'), allow_pickle=True).item()
            # ref = cv2.imread(os.path.join(self.template_path, barcode+'10.jpg'))
            # ref = cv2.cvtColor(ref, cv2.COLOR_BGR2RGB)
            # tar = self.image4
            # #finish0 = time.time()
            # #print('Read image time: ', finish0-start)
            # keypoints, descriptor, points = generateKeyPoint(attrs['keypoints']), attrs['descriptors'], attrs['position']
            # images = match(ref, tar, keypoints, descriptor, points)
            # labels10 = self.infer(images)

            # attrs = np.load(os.path.join(self.template_path, barcode+'12.npy'), allow_pickle=True).item()
            # ref = cv2.imread(os.path.join(self.template_path, barcode+'12.jpg'))
            # ref = cv2.cvtColor(ref, cv2.COLOR_BGR2RGB)
            # tar = self.image5
            # keypoints, descriptor, points = generateKeyPoint(attrs['keypoints']), attrs['descriptors'], attrs['position']
            # images = match(ref, tar, keypoints, descriptor, points)
            # labels12 = self.infer(images)
            # labels = labels10 + labels12
            # #finish = time.time()
            #print('Match cost: ', finish - start)

            # if np.argmax(labels) == 0:
            #     text = 'Black'
            # elif np.argmax(labels) == 1:
            #     text = 'Brown'
            # else:
            #     text = 'White'




            # 更改
            # self.image_zm = {
            #     "left": self.image3,
            #     "right": self.image4,
            #     "desc": "left or right :门板在图片的位置"
            # }
            # self.det_xd(self.cfg, self.image_zm ,barcode)
            # self.det_yt(self.cfg, self.image_zm ,barcode)
            # self.det_upr(self.cfg, self.image_zm ,barcode)
            # self.det_m(self.cfg, self.image_zm ,barcode)
            # self.det_lxzst(self.cfg, self.image_zm ,barcode)




            # self.det_lbz(self.image5 ,barcode)
            # self.det_ar(self.cfg, self.image_zm ,barcode)

            
        else:
            print("识识别出二维码... 无法进行面面检测")
        #     text = 'Unknow'

        # cv2.putText(self.image4, text, (1000, 500), cv2.FONT_HERSHEY_SIMPLEX, 5, (0,0,255), 5)
        # cv2.putText(self.image5, text, (1000, 500), cv2.FONT_HERSHEY_SIMPLEX, 5, (0,0,255), 5)


    def infer(self, images):
        labels = np.zeros(self.num_classes)
        for image in images:
            result, label = self.model_cls.predict(image)
            labels[int(label)] += 1
        return labels
        
    def toyota(self, data):
        time_start_get_imgs = time.time()
        self.get_imgs(data)
        time_start_deal_fm = time.time()

        barcode = self.deal_fm_attr()
        
        # 标签判别
        if barcode is not None:
            print(5*"**","门板类型:", barcode, 5*"**")
            # 样式判断
            self.style()
            
            # 检测误品
            # fm
            self.deal_fm(barcode)
            # zm
            time_start_deal_zm = time.time()
            self.deal_zm(barcode)
            time_end_deal_zm = time.time()
            # time
            print('get_imgs time: ', time_start_deal_fm-time_start_get_imgs)
            print('Det_fm time: ', time_start_deal_zm-time_start_deal_fm)
            print('Det_zm time: ', time_end_deal_zm-time_start_deal_zm)
            print('Total time : ', time_end_deal_zm-time_start_get_imgs)  
        else:
            print(5*"**","无标签，请正确粘贴标签", 5*"**")
        
        images = self.set_images()
        result = self.set_result()
        for k,v in result.items():
            print("{}:{}" .format(k, v) if v else '\033[1;31;40m{}:{}\033[0m'.format(k, v))
        # print('detection result:', result)
        return images, result

    def det_xd(self, cfg, image, barcode):
        # self.flag_xd = xd(cfg, image, barcode, visulization=self.cfg['plot_zm'], is_save=self.cfg['save_npy'])
        self.flag_xd = True
    def det_yt(self, cfg):
        # self.flag_yt = yt(cfg, image, barcode, visulization=self.cfg['plot_zm'], is_save=self.cfg['save_npy'])
        self.flag_yt = yt(cfg, self.Item_zm_left, self.Item_zm_right)
    def det_upr(self, cfg, image, barcode):
        # self.flag_upr = upr(cfg, image, barcode, visulization=self.cfg['plot_zm'], is_save=self.cfg['save_npy'])
        self.flag_upr = True
    def det_m(self, cfg, barcode):
        if 'R-R' in barcode: 
            m_flag = m(cfg, self.image4, self.Item_zm_right, self.shb_cls, barcode)
        else:
            m_flag = m(cfg, self.image3, self.Item_zm_left, self.shb_cls, barcode)
        self.flag_m = m_flag

    def det_lxzst(self, cfg):
        # self.flag_lxzst = lxzst(cfg, image, barcode, visulization=self.cfg['plot_zm'], is_save=self.cfg['save_npy'])
        self.flag_lxzst = lxzst(cfg, self.Item_zm_left, self.Item_zm_right)
    def det_lbz(self, cfg, image_zm, barcode):
        if 'RR-R' in barcode:
            self.flag_lbz = lbz(cfg, self.Item_zm_left, image_zm['left'], barcode, self.lbz_cls, visualization=True)
        elif 'RR-L' in barcode: 
            self.flag_lbz = lbz(cfg, self.Item_zm_right, image_zm['right'], barcode, self.lbz_cls, visualization=True)
        else:
            self.flag_lbz = True
    def det_ar(self, cfg, image, barcode):
        self.flag_ar = ar(cfg, image, barcode, visulization=self.cfg['plot_zm'], is_save=self.cfg['save_npy'])
    def det_zst(self, cfg, image, barcode):
        if 'FR-R' in barcode or 'FR-L' in barcode:
            self.flag_zst = zst(cfg, image, barcode, classif_net=self.zst_cls,  visulization=True)
        else:
            self.flag_zst = True


    def det_bq(self, Item1, Item2):
        if 'bq' in Item1.keys() or 'bq' in Item2.keys():
            self.flag_bq = True
        else:
            self.flag_bq = False
        
    def det_tp(self, cfg, item1, item2, barcode):
        tp2 = tp(cfg, item1, barcode)
        tp1 = tp(cfg, item2, barcode)
        if tp1 and tp2:
            self.flag_tp = True
        else:
            self.flag_tp = False
    def det_ls(self, cfg, item11, item13, barcode):
        # print(item11.keys(), item13.keys())
        if self.flag_tp:
            count1, self.image_ls8 = ls(cfg, self.image1, item11, barcode, "192.168.8.8", det_net=self.ls_cls, visualization=self.cfg['plot_ls'])
            count2, self.image_ls7 = ls(cfg, self.image2, item13, barcode, "192.168.7.7", det_net=self.ls_cls, visualization=self.cfg['plot_ls'])
            if count1 and count2:
                self.flag_ls = True
            else:
                self.flag_ls = False
        else:
            self.flag_ls = False

    def det_rd(self, cfg, item1, item2, barcode):
        if self.flag_tp:
            count1, self.image_rd8 = rd(cfg, self.image1, item1, barcode, "192.168.8.8", det_net=self.rd_cls, visualization=self.cfg['plot_rd'])
            count2, self.image_rd7 = rd(cfg, self.image2, item2, barcode, "192.168.7.7", det_net=self.rd_cls, visualization=self.cfg['plot_rd'])
            if count1 and count2:
                self.flag_rd = True
            else:
                self.flag_rd = False
        else:
            self.flag_rd = False
            return False
    def det_kk(self, cfg, item1, item2, barcode):
        if self.flag_tp:
            count1, self.image_kk8 = kk(cfg, self.image1, item1, barcode, "192.168.8.8", det_net=self.kk_cls, visualization=self.cfg['plot_kk'])
            count2, self.image_kk7 = kk(cfg, self.image2, item2, barcode, "192.168.7.7", det_net=self.kk_cls, visualization=self.cfg['plot_kk'])
            if count1 and count2 :
                self.flag_kk = True
            else:
                self.flag_kk = False
        else:
            self.flag_kk = False

    def det_lsh(self, cfg, item11, barcode):
        self.flag_lsh = lsh(cfg, item11, barcode)
    def det_xs(self, cfg, item13, barcode):
    
        self.flag_xs = xs(cfg, item13, barcode)
    def det_eapm(self, cfg, item11, item12, barcode):
        self.flag_eapm = eapm(cfg, item11, barcode) or eapm(cfg, item12, barcode)

    def det_mz(self, cfg):
        print('not implement ...')

    def det_mq(self):
        self.flag_mq = True
        
    def det_ewm(self):
        self.flag_ewm = True

    def det_bq_ocr(self, cfg, item1, item2):
        res1 = ocr_get_number(cfg, self.image1, item1, self.text_sys)
        res2 = ocr_get_number(cfg, self.image2, item2, self.text_sys)
        if res1:
            return res1
        if res2:
            return res2



def abc(data):
#if __name__ == '__main__':
#    img_path = 'TestData/C50_FR-R/171118'
#
#    data = dict()
#    for img_dir in os.listdir(img_path):
#        if not img_dir.endswith('jpg'):
#            continue
#        name = img_dir.strip('.jpg')
#        image = cv2.imread(os.path.join(img_path, img_dir))
#        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
#        data[name] = image
    print('start = time.time()')
    start = time.time()
    model = Toyota('/home/det/桌面/chengdu/chengdu/TensorRT/config.yaml')
    finish1 = time.time()
    images, result = model.toyota(data=data)
    finish2 = time.time()
    print('Total time : ', finish2-start)
    print('Det time: ', finish2-finish1)

    for key, value in images.items():
        #value = cv2.cvtColor(value, cv2.COLOR_RGB2BGR)
        cv2.imwrite('./TensorRT/Test'+'/'+key+'.jpg', value)

    print('over')



if __name__ == "__main__":
    test_folder = os.listdir('/media/Harddisk_B/Datasets/door_problem_detasets/chengdu/nodefect/')
    ind = random.randint(0, len(test_folder)-1)
    folder_name = test_folder[ind]
    folder_name = '124703_C50FR-R'
    img_path = os.path.join('/media/Harddisk_B/Datasets/door_problem_detasets/chengdu/nodefect', folder_name)

    data = dict()
    for img_dir in os.listdir(img_path):
        if not img_dir.endswith('jpg'):
            continue
        if img_dir.startswith('._'):
            continue
        if not img_dir.startswith('original_'):
            continue
        print("loading..", img_dir)
        
        name = img_dir.strip('.jpg')
        
        image = cv2.imread(os.path.join(img_path, img_dir)) # 输入是BGR OpenCV是BGR格式，PIL是RGB格式
        
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        data[name] = image

    start = time.time()
    model = Toyota('./config.yaml')
    finish1 = time.time()
    images, result = model.toyota(data=data)
    # print(result)
    finish2 = time.time()
    print('Total time : ', finish2-start)
    print('Det time: ', finish2-finish1)

    for key, value in images.items():
        #value = cv2.cvtColor(value, cv2.COLOR_RGB2BGR)
        cv2.imwrite('Test'+'/'+key+'.jpg', value)

    print('over')